import cv2
import numpy as np

class Mouse:
    def __init__(self, img):
        self.startX = 0
        self.startY = 0
        self.flag = False
        self.img = img
        self.copyImg = img.copy()
        self.mask = np.zeros(img.shape[:2], dtype=np.uint8)
        self.rect = (0, 0, 0, 0)
        self.output = np.zeros(img.shape, dtype=np.uint8)

    def onmouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            self.flag = True
            self.startX = x
            self.startY = y
            print("left down")
        elif event == cv2.EVENT_LBUTTONUP:
            self.flag = False
            cv2.rectangle(self.img, (self.startX, self.startY), (x, y), (255, 0, 0), 3)
            self.rect = (min(self.startX, x), min(self.startY, y), abs(self.startX - x), abs(self.startY - y))
            print("left up")
        elif event == cv2.EVENT_MOUSEMOVE:
            print(self.flag)
            if self.flag == True:
                self.img = self.copyImg.copy() 
                cv2.rectangle(self.img, (self.startX, self.startY), (x, y), (255, 0, 0), 3)
            print("move")

    def run(self):
        print("run")
        cv2.namedWindow('img', cv2.WINDOW_AUTOSIZE)
        cv2.setMouseCallback('img', self.onmouse)
        while True:
            cv2.imshow('img', self.img)
            cv2.imshow('output', self.output)
            key = cv2.waitKey(10)
            if key & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                break

            if key & 0xFF == ord('d'):
                bgdModel = np.zeros((1, 65), np.float64)
                fgdModel = np.zeros((1, 65), np.float64)
                cv2.grabCut(self.copyImg, self.mask, self.rect, bgdModel, fgdModel, 1, cv2.GC_INIT_WITH_RECT)

            if key == ord('s'):
                cv2.imwrite('./img/newsep.png', self.output)
                cv2.destroyAllWindows()
                break

            mask_t = np.where((self.mask == 1) | (self.mask == 3), 255, 0).astype('uint8')
            self.output = cv2.bitwise_and(self.copyImg, self.copyImg, mask=mask_t)
            
def separate(img):
    Mouse(img).run()
